import codecs

prompt_path = '/home/rickwwang/project_research/fairseq/data-bin/prompt.txt.test'
all_prompts = codecs.open(prompt_path, 'r', encoding='utf8').readlines()
all_prompts = map(lambda x: x.strip(), all_prompts)

dict_path = '/home/rickwwang/project_research/fairseq/data-bin/writingPrompts-prompt2srl2_500/dict.prompt.txt'
words = codecs.open(dict_path, 'r', encoding='utf8').readlines()
words = map(lambda x: x.strip().split(' ')[0], words)
prompt_dict = set(words)

exp_file_src = []
# exp_file_src.append('out/prompt2srl2_500/49src.txt')
# exp_file_src.append('out/prompt2srl2_500/49src.txt')
# exp_file_src.append('out/prompt2srl2_500/49src.txt')
# exp_file_src.append('out/prompt2srl2_500_wo_smooth_d1/40src.txt')
exp_file_src.append('fire_data/running/prompt2srl2_500_gpt2_nosympol/47src.txt')
exp_file_src.append('fire_data/running/prompt2srl2_500_gpt2_nosympol/47src.txt')
exp_file_src.append('fire_data/running/prompt2srl2_500_gpt2_nosympol/47src.txt')
exp_file_src.append('fire_data/running/prompt2srl2_500_gpt2_nosympol/47src.txt')
exp_file_src.append('fire_data/running/prompt2srl2_500_gpt2_nosympol/47src.txt')
exp_file_src.append('fire_data/running/prompt2srl2_500_gpt2_nosympol/47src.txt')
exp_file_src.append('fire_data/running/prompt2srl2_500_wo_smooth_d1_gpt2_nosympol/12src.txt')
exp_file_src.append('fire_data/running/prompt2srl2_500_wo_smooth_d1_gpt2_nosympol/12src.txt')
exp_file_src.append('fire_data/running/prompt2srl2_500_wo_smooth_d1_gpt2_nosympol/12src.txt')
exp_file_src.append('fire_data/running/prompt2srl2_500_wo_smooth_d1_gpt2_nosympol/12src.txt')
exp_file_src.append('fire_data/running/prompt2srl2_500_wo_smooth_d1_gpt2_nosympol/12src.txt')
exp_file_tgt = []
# exp_file_tgt.append('out/prompt2srl2_500/49tgt.txt')
# exp_file_tgt.append('out/prompt2srl2_500/49hyp.txt.1.topk10.t1.0.nr6')
# exp_file_tgt.append('out/prompt2srl2_500/49hyp.txt.1.topk20.t1.0.nr6')
# exp_file_tgt.append('out/prompt2srl2_500_wo_smooth_d1/40hyp.txt.1.topk20.t1.0.nr6')
# exp_file_tgt.append('out/prompt2srl2_500_wo_smooth_d1/40hyp.txt.1.topk10.t1.0.nr6')
exp_file_tgt.append('fire_data/running/prompt2srl2_500_gpt2_nosympol/47tgt.txt')
exp_file_tgt.append('fire_data/running/prompt2srl2_500_gpt2_nosympol/47hyp.txt.1.topk20.t1.0.nr6')
exp_file_tgt.append('fire_data/running/prompt2srl2_500_gpt2_nosympol/47hyp.txt.11.topk20.t0.8.nr6')
exp_file_tgt.append('fire_data/running/prompt2srl2_500_gpt2_nosympol/47hyp.txt.11.topk20.t0.6.nr6')
exp_file_tgt.append('fire_data/running/prompt2srl2_500_gpt2_nosympol/47hyp.txt.11.topk20.t0.4.nr6')
exp_file_tgt.append('fire_data/running/prompt2srl2_500_gpt2_nosympol/47hyp.txt.11.topk20.t0.2.nr6')
exp_file_tgt.append('fire_data/running/prompt2srl2_500_wo_smooth_d1_gpt2_nosympol/12hyp.txt.1.topk20.t1.0.nr6')
exp_file_tgt.append('fire_data/running/prompt2srl2_500_wo_smooth_d1_gpt2_nosympol/12hyp.txt.11.topk20.t0.8.nr6')
exp_file_tgt.append('fire_data/running/prompt2srl2_500_wo_smooth_d1_gpt2_nosympol/12hyp.txt.11.topk20.t0.6.nr6')
exp_file_tgt.append('fire_data/running/prompt2srl2_500_wo_smooth_d1_gpt2_nosympol/12hyp.txt.11.topk20.t0.4.nr6')
exp_file_tgt.append('fire_data/running/prompt2srl2_500_wo_smooth_d1_gpt2_nosympol/12hyp.txt.11.topk20.t0.3.nr6')

base_path = '/home/rickwwang/project_research/'

for exp1, exp2 in zip(exp_file_src, exp_file_tgt):
    story_prompt_path = base_path + exp1
    story_path = base_path + exp2
    result_path = story_path + '.human.test'

    all_stories_prompts = codecs.open(story_prompt_path, 'r', encoding='utf8').readlines()
    all_stories = codecs.open(story_path, 'r', encoding='utf8').readlines()
    all_stories = zip(all_stories_prompts, all_stories)

    promp_to_story = dict()
    for s in all_stories:
        key = s[0].strip()
        value = s[1].strip()
        if key in promp_to_story:
            promp_to_story[key].append(value)
        else:
            promp_to_story[key] = [value]

    format_stories = []
    for p in all_prompts:
        p = p.split(' ')
        p = [w if w in prompt_dict else '<unk>' for w in p]
        p = ' '.join(p)
        ss = promp_to_story[p]

        format_stories.append(p)
        format_stories.append(ss[0])

    with codecs.open(result_path, 'w', encoding='utf8') as fout:
        fout.write('\n'.join(format_stories))
        fout.write('\n')